#!/usr/bin/env python3
"""
MCP 协议服务器实现

本模块实现了 Model Context Protocol (MCP) 服务器的核心功能。
主要功能包括：
- MCP 协议消息处理
- 资源和工具的注册与管理
- 客户端请求路由
- 错误处理和响应格式化
- 服务器生命周期管理
"""

import asyncio
import json
import sys
from typing import Any, Dict, List, Optional, Callable, Union
from dataclasses import dataclass, asdict
from datetime import datetime

from .logger import get_logger, mcp_logger, log_execution_time
from .connection_pool import DatabaseConnectionPool
from .config import DatabaseConfig
from .utils import safe_json_dumps


# MCP 协议常量
MCP_VERSION = "2024-11-05"
SERVER_NAME = "mysql-mcp-server"
SERVER_VERSION = "1.0.0"


@dataclass
class MCPResource:
    """
    MCP 资源定义
    
    表示一个可以被客户端访问的资源，如数据库表、视图等。
    """
    uri: str
    name: str
    description: str
    mime_type: str = "application/json"
    metadata: Optional[Dict[str, Any]] = None


@dataclass
class MCPTool:
    """
    MCP 工具定义
    
    表示一个可以被客户端调用的工具，如 SQL 查询执行器。
    """
    name: str
    description: str
    input_schema: Dict[str, Any]
    handler: Callable


@dataclass
class MCPRequest:
    """
    MCP 请求消息
    """
    jsonrpc: str = "2.0"
    id: Optional[Union[str, int]] = None
    method: Optional[str] = None
    params: Optional[Dict[str, Any]] = None


@dataclass
class MCPResponse:
    """
    MCP 响应消息
    """
    jsonrpc: str = "2.0"
    id: Optional[Union[str, int]] = None
    result: Optional[Any] = None
    error: Optional[Dict[str, Any]] = None


@dataclass
class MCPError:
    """
    MCP 错误定义
    """
    code: int
    message: str
    data: Optional[Any] = None


class MCPProtocolError(Exception):
    """
    MCP 协议错误异常
    
    当 MCP 协议处理过程中发生错误时抛出此异常。
    """
    
    def __init__(self, message: str, code: int, data: Optional[Any] = None):
        """
        初始化 MCP 协议错误
        
        Args:
            message: 错误消息
            code: 错误代码
            data: 附加数据
        """
        self.message = message
        self.code = code
        self.data = data
        super().__init__(message)
    
    def __str__(self) -> str:
        """
        返回错误消息字符串表示
        
        Returns:
            str: 错误消息
        """
        return self.message
    
    def to_dict(self) -> Dict[str, Any]:
        """
        转换为字典格式
        
        Returns:
            Dict[str, Any]: 错误字典
        """
        error_dict = {
            "code": self.code,
            "message": self.message
        }
        if self.data is not None:
            error_dict["data"] = self.data
        return error_dict


class MCPResourceError(Exception):
    """
    MCP 资源错误异常
    
    当 MCP 资源操作发生错误时抛出此异常。
    """
    
    def __init__(self, message: str, resource_name: Optional[str] = None):
        """
        初始化 MCP 资源错误
        
        Args:
            message: 错误消息
            resource_name: 资源名称
        """
        self.message = message
        self.resource_name = resource_name
        super().__init__(message)


class MCPToolError(Exception):
    """
    MCP 工具错误异常
    
    当 MCP 工具操作发生错误时抛出此异常。
    """
    
    def __init__(self, message: str, tool_name: Optional[str] = None):
        """
        初始化 MCP 工具错误
        
        Args:
            message: 错误消息
            tool_name: 工具名称
        """
        self.message = message
        self.tool_name = tool_name
        super().__init__(message)


class MCPServer:
    """
    MCP 协议服务器
    
    实现 MCP 协议的服务器端，处理客户端请求并提供数据库访问功能。
    """
    
    # MCP 错误代码
    ERROR_PARSE_ERROR = -32700
    ERROR_INVALID_REQUEST = -32600
    ERROR_METHOD_NOT_FOUND = -32601
    ERROR_INVALID_PARAMS = -32602
    ERROR_INTERNAL_ERROR = -32603
    ERROR_SERVER_ERROR = -32000
    
    def __init__(self, db_config: DatabaseConfig):
        """
        初始化 MCP 服务器
        
        Args:
            db_config: 数据库配置
        """
        self.logger = get_logger(__name__)
        self.db_config = db_config
        self.db_pool: Optional[DatabaseConnectionPool] = None
        
        # 注册的资源和工具
        self.resources: Dict[str, MCPResource] = {}
        self.tools: Dict[str, MCPTool] = {}
        
        # 方法处理器映射
        self.method_handlers = {
            "initialize": self._handle_initialize,
            "resources/list": self._handle_list_resources,
            "resources/read": self._handle_read_resource,
            "tools/list": self._handle_list_tools,
            "tools/call": self._handle_call_tool,
            "ping": self._handle_ping,
        }
        
        # 服务器状态
        self.initialized = False
        self.client_info: Optional[Dict[str, Any]] = None
        
        self._register_default_resources()
        self._register_default_tools()
    
    def get_server_info(self) -> Dict[str, Any]:
        """获取服务器信息"""
        return {
            'name': getattr(self, 'server_name', 'mysql-mcp-server'),
            'version': getattr(self, 'server_version', '1.0.0'),
            'is_running': self.initialized,
            'resources_count': len(self.resources),
            'tools_count': len(self.tools),
            'database_config': {
                'host': self.db_config.host,
                'port': self.db_config.port,
                'database': self.db_config.database,
                'user': self.db_config.user
            }
        }
    
    async def start(self) -> None:
        """
        启动 MCP 服务器
        
        初始化数据库连接池并准备处理客户端请求。
        """
        try:
            # 初始化数据库连接池
            self.db_pool = DatabaseConnectionPool(self.db_config)
            await self.db_pool.initialize()
            
            mcp_logger.log_server_event(
                "startup",
                server_name=SERVER_NAME,
                server_version=SERVER_VERSION,
                mcp_version=MCP_VERSION,
                database_host=self.db_config.host,
                database_port=self.db_config.port,
                database_name=self.db_config.database
            )
            
            self.logger.info(
                "MCP服务器启动成功",
                server_name=SERVER_NAME,
                version=SERVER_VERSION,
                database=f"{self.db_config.host}:{self.db_config.port}/{self.db_config.database}"
            )
            
        except Exception as e:
            self.logger.error("MCP服务器启动失败", error=str(e))
            mcp_logger.log_server_event("startup_error", error=str(e))
            raise
    
    async def stop(self) -> None:
        """
        停止 MCP 服务器
        
        清理资源并关闭数据库连接。
        """
        try:
            if self.db_pool:
                self.db_pool.close()
            
            mcp_logger.log_server_event("shutdown")
            self.logger.info("MCP服务器已停止")
            
        except Exception as e:
            self.logger.error("MCP服务器停止时发生错误", error=str(e))
            mcp_logger.log_server_event("shutdown_error", error=str(e))
    
    @log_execution_time()
    async def handle_request(self, request_data: str) -> str:
        """
        处理客户端请求
        
        Args:
            request_data: JSON-RPC 请求数据
            
        Returns:
            str: JSON-RPC 响应数据
        """
        try:
            # 解析请求
            request_dict = json.loads(request_data)
            request = MCPRequest(**request_dict)
            
            mcp_logger.log_request(
                method=request.method,
                params=request.params,
                request_id=request.id
            )
            
            # 处理请求
            response = await self._process_request(request)
            
            # 记录响应
            mcp_logger.log_response(
                method=request.method,
                success=response.error is None,
                result=response.result,
                error=response.error.get("message") if response.error else None,
                request_id=request.id
            )
            
            # 序列化响应，成功时不包含error字段
            response_dict = asdict(response)
            if response_dict.get('error') is None:
                response_dict.pop('error', None)
            return safe_json_dumps(response_dict, ensure_ascii=False)
            
        except json.JSONDecodeError as e:
            self.logger.error("JSON解析错误", error=str(e))
            error_response = MCPResponse(
                id=None,
                error=asdict(MCPError(
                    code=self.ERROR_PARSE_ERROR,
                    message="Parse error",
                    data=str(e)
                ))
            )
            return safe_json_dumps(asdict(error_response), ensure_ascii=False)
        
        except Exception as e:
            self.logger.error("处理请求时发生未知错误", error=str(e))
            error_response = MCPResponse(
                id=getattr(request, 'id', None) if 'request' in locals() else None,
                error=asdict(MCPError(
                    code=self.ERROR_INTERNAL_ERROR,
                    message="Internal error",
                    data=str(e)
                ))
            )
            return safe_json_dumps(asdict(error_response), ensure_ascii=False)
    
    async def _process_request(self, request: MCPRequest) -> MCPResponse:
        """
        处理单个请求
        
        Args:
            request: MCP 请求对象
            
        Returns:
            MCPResponse: MCP 响应对象
        """
        if not request.method:
            return MCPResponse(
                id=request.id,
                error=asdict(MCPError(
                    code=self.ERROR_INVALID_REQUEST,
                    message="Missing method"
                ))
            )
        
        # 查找方法处理器
        handler = self.method_handlers.get(request.method)
        if not handler:
            return MCPResponse(
                id=request.id,
                error=asdict(MCPError(
                    code=self.ERROR_METHOD_NOT_FOUND,
                    message=f"Method not found: {request.method}"
                ))
            )
        
        try:
            # 调用处理器
            result = await handler(request.params or {})
            return MCPResponse(id=request.id, result=result)
            
        except ValueError as e:
            return MCPResponse(
                id=request.id,
                error=asdict(MCPError(
                    code=self.ERROR_INVALID_PARAMS,
                    message=str(e)
                ))
            )
        
        except Exception as e:
            self.logger.error(f"处理方法 {request.method} 时发生错误", error=str(e))
            return MCPResponse(
                id=request.id,
                error=asdict(MCPError(
                    code=self.ERROR_SERVER_ERROR,
                    message=f"Server error: {str(e)}"
                ))
            )
    
    async def _handle_initialize(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理初始化请求
        
        Args:
            params: 初始化参数
            
        Returns:
            Dict[str, Any]: 服务器能力信息
        """
        self.client_info = params.get("clientInfo", {})
        self.initialized = True
        
        self.logger.info(
            "客户端已连接",
            client_name=self.client_info.get("name"),
            client_version=self.client_info.get("version")
        )
        
        return {
            "protocolVersion": MCP_VERSION,
            "serverInfo": {
                "name": SERVER_NAME,
                "version": SERVER_VERSION
            },
            "capabilities": {
                "resources": {
                    "subscribe": False,
                    "listChanged": False
                },
                "tools": {
                    "listChanged": False
                }
            }
        }
    
    async def _handle_list_resources(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理资源列表请求
        
        Args:
            params: 请求参数
            
        Returns:
            Dict[str, Any]: 资源列表
        """
        if not self.initialized:
            raise ValueError("Server not initialized")
        
        # 动态获取数据库表信息
        await self._update_table_resources()
        
        resources_list = [
            {
                "uri": resource.uri,
                "name": resource.name,
                "description": resource.description,
                "mimeType": resource.mime_type
            }
            for resource in self.resources.values()
        ]
        
        return {"resources": resources_list}
    
    async def _handle_read_resource(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理资源读取请求
        
        Args:
            params: 请求参数，包含资源 URI
            
        Returns:
            Dict[str, Any]: 资源内容
        """
        if not self.initialized:
            raise ValueError("Server not initialized")
        
        uri = params.get("uri")
        if not uri:
            raise ValueError("Missing uri parameter")
        
        resource = self.resources.get(uri)
        if not resource:
            raise ValueError(f"Resource not found: {uri}")
        
        # 根据资源类型返回不同内容
        if uri.startswith("mysql://table/"):
            table_name = uri.split("/")[-1]
            return await self._get_table_schema(table_name)
        
        return {
            "contents": [{
                "uri": resource.uri,
                "mimeType": resource.mime_type,
                "text": safe_json_dumps(asdict(resource), ensure_ascii=False, indent=2)
            }]
        }
    
    async def _handle_list_tools(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理工具列表请求
        
        Args:
            params: 请求参数
            
        Returns:
            Dict[str, Any]: 工具列表
        """
        if not self.initialized:
            raise ValueError("Server not initialized")
        
        tools_list = [
            {
                "name": tool.name,
                "description": tool.description,
                "inputSchema": tool.input_schema
            }
            for tool in self.tools.values()
        ]
        
        return {"tools": tools_list}
    
    async def _handle_call_tool(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理工具调用请求
        
        Args:
            params: 请求参数，包含工具名称和参数
            
        Returns:
            Dict[str, Any]: 工具执行结果
        """
        if not self.initialized:
            raise ValueError("Server not initialized")
        
        tool_name = params.get("name")
        if not tool_name:
            raise ValueError("Missing tool name")
        
        tool = self.tools.get(tool_name)
        if not tool:
            raise ValueError(f"Tool not found: {tool_name}")
        
        tool_params = params.get("arguments", {})
        
        # 调用工具处理器
        result = await tool.handler(tool_params)
        
        return {
            "content": [{
                "type": "text",
                "text": safe_json_dumps(result, ensure_ascii=False, indent=2)
            }]
        }
    
    async def _handle_ping(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        处理 ping 请求
        
        Args:
            params: 请求参数
            
        Returns:
            Dict[str, Any]: pong 响应
        """
        return {"status": "pong", "timestamp": datetime.now().isoformat()}
    
    def _register_default_resources(self) -> None:
        """
        注册默认资源
        
        注册服务器信息等静态资源。
        """
        # 服务器信息资源
        self.resources["mysql://server/info"] = MCPResource(
            uri="mysql://server/info",
            name="服务器信息",
            description="MySQL MCP 服务器的基本信息和配置",
            mime_type="application/json"
        )
    
    def _register_default_tools(self) -> None:
        """
        注册默认工具
        
        注册 SQL 查询执行器等核心工具。
        """
        # SQL 查询工具
        self.tools["execute_query"] = MCPTool(
            name="execute_query",
            description="执行 SQL 查询并返回结果",
            input_schema={
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "要执行的 SQL 查询语句"
                    },
                    "limit": {
                        "type": "integer",
                        "description": "结果限制数量（可选）",
                        "minimum": 1,
                        "maximum": 10000,
                        "default": 1000
                    }
                },
                "required": ["query"]
            },
            handler=self._execute_sql_query
        )
        
        # 表结构查询工具
        self.tools["describe_table"] = MCPTool(
            name="describe_table",
            description="获取数据库表的结构信息",
            input_schema={
                "type": "object",
                "properties": {
                    "table_name": {
                        "type": "string",
                        "description": "表名称"
                    }
                },
                "required": ["table_name"]
            },
            handler=self._describe_table
        )
        
        # 列出表工具
        self.tools["list_tables"] = MCPTool(
            name="list_tables",
            description="列出数据库中的所有表",
            input_schema={
                "type": "object",
                "properties": {},
                "required": []
            },
            handler=self._list_tables
        )
    
    async def _update_table_resources(self) -> None:
        """
        更新数据库表资源
        
        动态获取数据库中的表信息并注册为资源。
        """
        if not self.db_pool:
            return
        
        try:
            # 获取所有表名
            tables = self.db_pool.execute_query(
                "SHOW TABLES"
            )
            
            # 注册表资源
            for table_row in tables:
                table_name = list(table_row.values())[0]  # 获取表名
                uri = f"mysql://table/{table_name}"
                
                self.resources[uri] = MCPResource(
                    uri=uri,
                    name=f"表: {table_name}",
                    description=f"MySQL 数据库表 {table_name} 的结构和数据",
                    mime_type="application/json"
                )
                
        except Exception as e:
            self.logger.warning("更新表资源时发生错误", error=str(e))
    
    async def _get_table_schema(self, table_name: str) -> Dict[str, Any]:
        """
        获取表结构信息
        
        Args:
            table_name: 表名
            
        Returns:
            Dict[str, Any]: 表结构信息
        """
        if not self.db_pool:
            raise ValueError("Database connection not available")
        
        try:
            # 获取表结构，使用反引号包围表名以避免关键字冲突
            columns = self.db_pool.execute_query(
                f"DESCRIBE `{table_name}`"
            )
            
            return {
                "contents": [{
                    "uri": f"mysql://table/{table_name}",
                    "mimeType": "application/json",
                    "text": safe_json_dumps({
                        "table_name": table_name,
                        "columns": columns,
                        "timestamp": datetime.now().isoformat()
                    }, ensure_ascii=False, indent=2)
                }]
            }
            
        except Exception as e:
            self.logger.error(f"获取表 {table_name} 结构时发生错误", error=str(e))
            # 返回错误信息而不是抛出异常
            return {
                "contents": [{
                    "uri": f"mysql://table/{table_name}",
                    "mimeType": "application/json",
                    "text": safe_json_dumps({
                        "table_name": table_name,
                        "error": f"无法获取表结构: {str(e)}",
                        "error_type": type(e).__name__,
                        "timestamp": datetime.now().isoformat()
                    }, ensure_ascii=False, indent=2)
                }]
            }
    
    async def _execute_sql_query(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        执行 SQL 查询工具处理器
        
        Args:
            params: 工具参数
            
        Returns:
            Dict[str, Any]: 查询结果
        """
        if not self.db_pool:
            raise ValueError("Database connection not available")
        
        query = params.get("query")
        limit = params.get("limit", 1000)
        
        if not query:
            raise ValueError("Missing query parameter")
        
        try:
            # 执行查询
            result = self.db_pool.execute_query(
                query
            )
            
            return {
                "query": query,
                "result_count": len(result),
                "results": result,
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error("执行SQL查询时发生错误", query=query, error=str(e))
            raise
    
    async def _describe_table(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        描述表结构工具处理器
        
        Args:
            params: 工具参数
            
        Returns:
            Dict[str, Any]: 表结构信息
        """
        table_name = params.get("table_name")
        if not table_name:
            raise ValueError("Missing table_name parameter")
        
        if not self.db_pool:
            raise ValueError("Database connection not available")
        
        try:
            # 直接执行DESCRIBE查询，不使用_get_table_schema方法
            columns = self.db_pool.execute_query(f"DESCRIBE `{table_name}`")
            
            return {
                "table_name": table_name,
                "columns": columns,
                "count": len(columns),
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error(f"获取表 {table_name} 结构时发生错误", error=str(e))
            return {
                "table_name": table_name,
                "error": f"无法获取表结构: {str(e)}",
                "error_type": type(e).__name__,
                "timestamp": datetime.now().isoformat()
            }
    
    async def _list_tables(self, params: Dict[str, Any]) -> Dict[str, Any]:
        """
        列出所有表工具处理器
        
        Args:
            params: 工具参数
            
        Returns:
            Dict[str, Any]: 表列表信息
        """
        if not self.db_pool:
            raise ValueError("Database connection not available")
        
        try:
            # 获取所有表名
            tables = self.db_pool.execute_query("SHOW TABLES")
            
            # 提取表名列表
            table_names = []
            for table_row in tables:
                table_name = list(table_row.values())[0]  # 获取表名
                table_names.append(table_name)
            
            return {
                "tables": table_names,
                "count": len(table_names),
                "timestamp": datetime.now().isoformat()
            }
            
        except Exception as e:
            self.logger.error("列出表时发生错误", error=str(e))
            raise


async def create_mcp_server(db_config: DatabaseConfig) -> MCPServer:
    """
    创建并启动 MCP 服务器
    
    Args:
        db_config: 数据库配置
        
    Returns:
        MCPServer: 已启动的 MCP 服务器实例
    """
    server = MCPServer(db_config)
    await server.start()
    return server